#include <iostream>
#include <map>
#include <fstream>      // For file operations
#include <filesystem>   // For std::filesystem

/* 命令行参数解析*/
#include <CLI/CLI.hpp>

// GUI 相关的头文件
#include <QApplication>
#include "GUI.h"

#include "signature_manager/RSA.h"
#include "signature_manager/DSA.h"
#include "signature_manager/Rabin.h"
#include "signature_manager/ECDSA.h"
#include "signature_manager/EdDSA.h"
#include "signature_manager/Hash.h"
#include "signature_manager/signature_manager.h"
#include "signature_manager/data_conversion.h"
#include "signature_manager/GMPInt.h"

#include "DEBUG_PRINT_MACRO.h"


int main(int argc, char* argv []) {

    #ifdef DEBUG_PRINTS_LEVEL_1
    std::cout << "[main][info]Debug prints are enabled.\n";
    #endif

        // 检查是否存在 "gui" 命令行参数
    bool runGui = false;
    for (int i = 1; i < argc; ++i) {
        if (std::string(argv[i]) == "--gui") {
            runGui = true;
            // 移除 --gui 参数，因为它不是 CLI11 或 Qt 应用程序的常规参数
            // 创建一个新的 argv 数组，不包含 --gui
            char** new_argv = new char*[argc];
            int new_argc = 0;
            for (int j = 0; j < argc; ++j) {
                if (j == i) continue; // Skip --gui
                new_argv[new_argc++] = argv[j];
            }
            argc = new_argc;
            argv = new_argv;
            break;
        }
    }

    if (runGui) {
        QApplication app(argc, argv);
        SignatureGUI gui;
        gui.show();
        return app.exec();
    }


    CLI::App app{ "Signature Tool" };
    if (argc < 2) {
        std::cout << app.help() << std::endl;
        return 0;
    }

    std::string input_file;            // 待签名/验签的文件路径
    std::string signature_algorithm;   // 签名算法
    int key_length;                    // 密钥长度
    std::string curve_name;            // 曲线名称 (ECDSA/EdDSA)


    std::string public_parameters;     // 公共参数 (命令行输入)
    std::string public_key;            // 公钥 (命令行输入)
    std::string signature_str;         // 签名字符串 (命令行输入)
    std::string sign_output_file;      // 签名输出文件路径 (sign 子命令)
    std::string sign_input_file;       // 签名输入文件路径 (verify 子命令)


    // 子命令 "sign" 签名
    auto sign = app.add_subcommand("sign", "对文件进行签名");
    sign->add_option("-f,--file", input_file, "待签名的文件")->required();
    sign->add_option("-a,--algorithm", signature_algorithm, "签名算法, 如 RSA DSA Rabin ECDSA EdDSA")->required();
    sign->add_option("-l,--length", key_length, "密钥长度, 例如 1024, 2048, 3072")->default_val(1024); // 默认值改为整数
    sign->add_option("-c,--curve", curve_name, "曲线名称, 如 secp256k1, prime256v1");
    sign->add_option("-o,--output", sign_output_file, "输出文件路径，用于保存签名、公共参数和公钥")->default_val("");

    // 子命令 "verify" 验签
    auto verify = app.add_subcommand("verify", "验证文件签名");
    verify->add_option("-f,--file", input_file, "待验证的文件")->required();
    verify->add_option("-a,--algorithm", signature_algorithm, "签名算法, 如 RSA DSA Rabin ECDSA EdDSA")->required();
    verify->add_option("-i,--input-signature-file", sign_input_file, "包含签名、公共参数和公钥的输入文件");
    verify->add_option("-p,--parameters", public_parameters, "公共参数，命令行输入");
    verify->add_option("-k,--key", public_key, "公钥，命令行输入");
    verify->add_option("-s,--signature", signature_str, "签名字符串，命令行输入");

    try {
        app.parse(argc, argv);
    }
    catch (const CLI::ParseError& e) {
        return app.exit(e); // 自动处理错误和帮助信息
    }


    // 添加 CLI11 回调函数，用于验证参数的互斥性
    verify->preparse_callback([&] (size_t) {
        if (!sign_input_file.empty()) { // 如果提供了输入签名文件
            if (!public_key.empty() || !signature_str.empty() || !public_parameters.empty()) {
                throw CLI::ValidationError("当使用 -i 参数时，不能同时提供 -k, -s 和 -p 参数.");
            }
        }
        else { // 如果没有提供输入签名文件，则必须提供公钥和签名字符串
            if (public_key.empty() || signature_str.empty()) {
                throw CLI::ValidationError("对于验签，必须提供 -i 参数，或者同时提供 -k 和 -s 参数.");
            }
        }
    });



    PlainMessage sign_target_data(new SHA256()); // 待签名/验签的数据
    if (!input_file.empty()) { // 检查输入文件是否为空
        try {
            sign_target_data.read_file_to_bytes(input_file.c_str());
        }
        catch (const std::exception& e) {
            std::cerr << "错误：无法读取文件: " << e.what() << std::endl;
            return 1; // 错误：无法读取文件
        }
    }
    else {
        // 如果 input_file 为空，但是是 verify 子命令，这可能是预期的行为（例如只验证已知字符串的签名）
        // 对于本程序逻辑，input_file 总是 required，所以此处分支理论上不会被触发，但保留以防万一
        std::cerr << "未指定输入文件\n";
        return 1;
    }

    SignatureManager_base* signature_manager; // 签名管理器基类指针

    // 根据选择的签名算法和参数创建对应的签名管理器实例
    if (!signature_algorithm.empty()) {
        if (signature_algorithm == "RSA") {
            if (key_length == 1024) {
                signature_manager = new SignatureManager<RSA_type_traits<1024>>();
            }
            else if (key_length == 2048) {
                signature_manager = new SignatureManager<RSA_type_traits<2048>>();
            }
            else if (key_length == 3072) {
                signature_manager = new SignatureManager<RSA_type_traits<3072>>();
            }
            else {
                printf("RSA 密钥长度无效: %d\n", key_length);
                return 1; // 错误：无效的密钥长度
            }
        }
        else if (signature_algorithm == "DSA") {
            if (key_length == 1024) {
                signature_manager = new SignatureManager<DSA_type_traits<1024>>();
            }
            else if (key_length == 2048) {
                signature_manager = new SignatureManager<DSA_type_traits<2048>>();
            }
            else if (key_length == 3072) {
                signature_manager = new SignatureManager<DSA_type_traits<3072>>();
            }
            else {
                printf("DSA 密钥长度无效: %d\n", key_length);
                return 1; // 错误：无效的密钥长度
            }
        }
        else if (signature_algorithm == "Rabin") {
            if (key_length == 1024) {
                signature_manager = new SignatureManager<Rabin_type_traits<1024>>();
            }
            else if (key_length == 2048) {
                signature_manager = new SignatureManager<Rabin_type_traits<2048>>();
            }
            else if (key_length == 3072) {
                signature_manager = new SignatureManager<Rabin_type_traits<3072>>();
            }
            else {
                printf("Rabin 密钥长度无效: %d\n", key_length);
                return 1; // 错误：无效的密钥长度
            }
        }
        else if (signature_algorithm == "ECDSA") {
            if (curve_name == "secp256k1" || curve_name.empty()) { // 默认曲线
                signature_manager = new SignatureManager<ECDSA_type_traits<secp256k1_tag>>();
            }
            else if (curve_name == "prime256v1") {
                signature_manager = new SignatureManager<ECDSA_type_traits<prime256v1_tag>>();
            }
            else if (curve_name == "secp384r1") {
                signature_manager = new SignatureManager<ECDSA_type_traits<secp384r1_tag>>();
            }
            else {
                printf("ECDSA 曲线名称无效: %s\n", curve_name.c_str());
                return 1; // 错误：无效的曲线名称
            }
        }
        else if (signature_algorithm == "EdDSA") {
            if (curve_name == "Ed25519" || curve_name.empty()) { // 默认曲线
                signature_manager = new SignatureManager<EdDSA_type_traits<Ed25519_tag>>();
            }
            else if (curve_name == "Ed448") {
                signature_manager = new SignatureManager<EdDSA_type_traits<Ed448_tag>>();
            }
            else {
                printf("EdDSA 曲线名称无效: %s\n", curve_name.c_str());
                return 1; // 错误：无效的曲线名称
            }
        }
        else {
            std::cerr << "不支持的签名算法: " << signature_algorithm << std::endl;
            return 1;
        }
    }
    else {
        std::cerr << "未指定签名算法\n";
        return 1; // 错误：未指定签名算法
    }

    // 处理 "sign" 子命令
    if (sign->parsed()) {
        signature_manager->generateKeyPair(); // 生成密钥对

        std::string generated_signature_str = "";
        signature_manager->signMessage(sign_target_data, generated_signature_str); // 签名消息

        if (!sign_output_file.empty()) { // 如果指定了输出文件
            std::filesystem::path outputPath(sign_output_file);
            std::ofstream outFile(outputPath); // 打开输出文件
            if (outFile.is_open()) {
                // 将签名、公共参数和公钥以 Base64 编码形式写入文件，采用新格式
                outFile << "---PUBLIC KEY---\n";
                outFile << base64_encode(signature_manager->get_publickey()) << "\n";
                outFile << "---PUBLIC PARAMETERS---\n";
                outFile << base64_encode(signature_manager->get_publicparameters()) << "\n";
                outFile << "---SIGNATURE---\n";
                outFile << base64_encode(generated_signature_str) << "\n";
                outFile.close();
                std::cout << "签名和密钥已保存至: " << sign_output_file << std::endl;
            }
            else {
                std::cerr << "无法打开输出文件: " << sign_output_file << std::endl;
                delete signature_manager;
                return 1;
            }
        }
        else { // 如果没有指定输出文件，则打印到标准输出
            std::cout << "---PUBLIC KEY---\n";
            std::cout << base64_encode(signature_manager->get_publickey()) << std::endl;
            std::cout << "---PUBLIC PARAMETERS---\n";
            std::cout << base64_encode(signature_manager->get_publicparameters()) << std::endl;
            std::cout << "---SIGNATURE---\n";
            std::cout << base64_encode(generated_signature_str) << std::endl;
        }
    }
    // 处理 "verify" 子命令
    else if (verify->parsed()) {
        std::string decoded_public_key;
        std::string decoded_public_parameters;
        std::string decoded_signature_str;

        if (!sign_input_file.empty()) { // 如果提供了输入签名文件
            std::filesystem::path inputPath(sign_input_file);
            std::ifstream inFile(inputPath); // 打开输入文件
            if (!inFile.is_open()) {
                std::cerr << "无法打开输入签名文件: " << sign_input_file << std::endl;
                delete signature_manager;
                return 1;
            }

            std::string line;
            std::string current_section; // 用于跟踪当前读取的节
            while (std::getline(inFile, line)) {
                // 移除行首尾的空白字符
                line.erase(0, line.find_first_not_of(" \t\n\r\f\v"));
                line.erase(line.find_last_not_of(" \t\n\r\f\v") + 1);

                if (line.empty()) {
                    continue; // 跳过空行
                }
                
                if (line == "---PUBLIC KEY---") {
                    current_section = "PUBLIC_KEY";
                } else if (line == "---PUBLIC PARAMETERS---") {
                    current_section = "PUBLIC_PARAMETERS";
                } else if (line == "---SIGNATURE---") {
                    current_section = "SIGNATURE";
                } else {
                    // 根据当前节存储数据
                    if (current_section == "PUBLIC_KEY") {
                        public_key = line;
                    } else if (current_section == "PUBLIC_PARAMETERS") {
                        public_parameters = line;
                    } else if (current_section == "SIGNATURE") {
                        signature_str = line;
                    }
                }
            }
            inFile.close();

            // 检查是否所有必需字段都已解析
            if (public_key.empty() || signature_str.empty()) { // PublicParameters 可能是可选的，因此不强制检查
                std::cerr << "输入签名文件缺少必需字段: PUBLIC KEY 或 SIGNATURE.\n";
                delete signature_manager;
                return 1;
            }

            decoded_signature_str = base64_decode(signature_str);
            decoded_public_parameters = base64_decode(public_parameters);
            decoded_public_key = base64_decode(public_key);
        }
        else { // 如果没有提供输入签名文件，则使用命令行参数
            decoded_public_key = base64_decode(public_key);
            // public_parameters 可能是可选的，即使提供了，也可能为空，DSA等算法需要
            decoded_public_parameters = base64_decode(public_parameters);
            decoded_signature_str = base64_decode(signature_str);
        }

        // 设置签名管理器的公钥和公共参数
        signature_manager->set_publickey(decoded_public_key);
        signature_manager->set_publicparameters(decoded_public_parameters); // 即使为空，也调用以确保状态正确

        // 验证签名
        bool result = signature_manager->verifySignature(sign_target_data, decoded_signature_str);

        std::cout << "验签结果: " << (result ? "签名验证成功" : "签名验证失败") << std::endl;
    }

    delete signature_manager; // 释放内存
    return 0;
}
